from datetime import datetime
from typing import Dict, Any, List, TypedDict, AsyncGenerator

from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, END
from langgraph.graph.state import CompiledStateGraph
from langgraph.types import interrupt

from AIAgents.HumanInputLoopAgent import HumanInputLoopAgent
from AIAgents.actions.execute import execute_tools_node, execute_mcp_node, execute_rag_node
from AIAgents.actions.think import think_node, should_continue_execution
from config.logging_config import get_logger
from prompt.navigation_prompt import NavigationPrompts
from utils.SendResponseDTO import SendResponseDTO
from utils.loggerUtils import LoggerUtils

# 设置日志记录器
logger = get_logger(__name__)


class NavigationAgentState(TypedDict):
    """导航Agent的状态定义"""
    # 基础输入
    input: str
    original_input: str

    # 记忆上下文
    history_memory_context: Dict[str, Any]  # 历史记忆
    this_memory_context: List[Dict[str, Any]]  # 当前执行记忆

    # 思考过程
    current_thought: str
    reasoning_steps: List[Dict[str, Any]]

    # 智能响应解析
    intelligent_response: Dict[str, Any]  # 存储LLM的智能响应
    parsed_action: str  # 解析后的action：'use_tool', 'use_mcp', 'use_rag', 'navigation', 'direct_answer', 'finish'
    action_params: Dict[str, Any]  # action的参数

    # 执行结果
    tool_results: List[Dict[str, Any]]
    mcp_results: List[Dict[str, Any]]
    rag_results: List[Dict[str, Any]]

    # 人机交互
    need_human_interaction: bool  # 是否需要人机交互
    human_interaction_message: Dict[str, Any]  # 人机交互消息
    user_confirmation: bool  # 用户确认结果

    # 状态控制
    next_action: str  # 'navigation', 'tools', 'mcp', 'rag', 'human', 'check_human', 'finish'
    should_continue: bool
    iteration_count: int
    max_iterations: int

    # 动态的mcp_list
    mcp_list: List[Dict[str, Any]]  # 可用的MCP服务列表
    # 最终输出
    final_response: str
    success: bool
    error: str


class NavigationAgent(HumanInputLoopAgent):

    def __init__(self, **kwargs):
        logger.info("🚀 初始化NavigationAgent")

        # 初始化NavigationAgent特有的属性
        self.current_config = None
        self.checkpointer = MemorySaver()

        super().__init__(**kwargs)

        self.log_component_status()

    def _build_graph(self) -> CompiledStateGraph:
        """构建简化的思考-执行循环工作流"""
        workflow = StateGraph(NavigationAgentState)

        # 定义节点
        self._define_nodes(workflow)

        # 定义边连接关系
        self._define_edges(workflow)

        # 编译图，使用checkpointer支持interrupt
        return workflow.compile(checkpointer=self.checkpointer)

    def _define_nodes(self, workflow: StateGraph) -> None:
        """定义工作流节点"""
        # 添加所有节点
        workflow.add_node("think", lambda state: think_node(self, state))
        workflow.add_node("check_human", self._check_human_node)
        workflow.add_node("human", self._interrupt_human_node)
        workflow.add_node("execute_tools", lambda state: execute_tools_node(self, state))
        workflow.add_node("execute_mcp", lambda state: execute_mcp_node(self, state))
        workflow.add_node("execute_rag", lambda state: execute_rag_node(self, state))
        workflow.add_node("finish", self._finish_node)

    def _define_edges(self, workflow: StateGraph) -> None:
        """定义工作流边连接关系"""
        # 设置入口点
        workflow.set_entry_point("think")

        # 添加条件边：思考后决定执行哪个节点
        workflow.add_conditional_edges(
            "think",
            self._decide_next_action,
            {
                "check_human": "check_human",
                "rag": "execute_rag",
                "finish": "finish"
            }
        )

        # 检查人机交互节点的条件边
        workflow.add_conditional_edges(
            "check_human",
            self._decide_after_human_check,
            {
                "human": "human",
                "tools": "execute_tools",
                "mcp": "execute_mcp"
            }
        )

        # 人机交互节点完成后进入实际执行节点
        workflow.add_conditional_edges(
            "human",
            self._decide_after_human_interaction,
            {
                "tools": "execute_tools",
                "mcp": "execute_mcp",
                "finish": "finish"
            }
        )

        # 执行节点完成后回到思考节点（除了finish）
        workflow.add_conditional_edges(
            "execute_tools",
            self._should_continue_thinking,
            {"continue": "think", "finish": "finish"}
        )

        workflow.add_conditional_edges(
            "execute_mcp",
            self._should_continue_thinking,
            {"continue": "think", "finish": "finish"}
        )

        workflow.add_conditional_edges(
            "execute_rag",
            self._should_continue_thinking,
            {"continue": "think", "finish": "finish"}
        )

        # finish节点的条件边：如果用户有更多问题则继续思考，否则结束
        workflow.add_conditional_edges(
            "finish",
            self._should_continue_after_finish,
            {"continue": "think", "end": END}
        )

    def _create_initial_state(self, input_text: str) -> Dict[str, Any]:
        """创建初始状态 - NavigationAgent特有实现"""
        # 使用父类的基础状态
        state = self._create_base_initial_state(input_text)

        # 准备包含记忆上下文的输入
        history_memory_context = self._prepare_input_with_memory(input_text)

        # 添加NavigationAgent特有的状态字段
        state.update({
            "history_memory_context": history_memory_context,
            "this_memory_context": [],  # 当前执行的记忆序列
            "current_thought": "",
            "reasoning_steps": [],
            "intelligent_response": {},
            "parsed_action": "",
            "tool_results": [],
            "mcp_results": [],
            "rag_results": [],
            "need_human_interaction": False,  # 是否需要人机交互
            "human_interaction_message": {},  # 人机交互消息
            "user_confirmation": False,  # 用户确认结果
            "next_action": ""
        })

        return state

    async def process_stream(self, input_data: Dict[str, Any]) -> AsyncGenerator[Dict[str, Any], None]:
        """流式处理导航请求"""
        try:
            input_text = input_data.get("input", "")
            LoggerUtils.log_agent_start("NavigationAgent", input_text, logger)

            if not input_text.strip():
                yield await self._yield_error_message("输入内容不能为空")
                return

            if not self.graph:
                yield await self._yield_error_message("导航工作流图未初始化")
                return

            yield await self._yield_thinking_message("导航Agent开始分析...", "start")

            # 初始化状态
            initial_state = self._create_initial_state(input_text)

            # 生成配置ID用于中断支持
            config = self._create_thread_config("nav_thread")
            self.current_config = config

            # 使用通用的流式处理方法
            graph_stream = self.graph.astream(initial_state, config=config, stream_mode="updates")
            async for result in self._process_graph_stream(graph_stream, config):
                yield result
        except Exception as e:
            LoggerUtils.log_error(f"NavigationAgent处理失败: {str(e)}", e, logger)
            yield await self._yield_error_message(f"导航处理错误: {str(e)}")

    def _finish_node(self, state: NavigationAgentState) -> NavigationAgentState:
        """结束节点 - 在任务完成时询问用户是否有其他补充问题"""
        logger.info("🏁 进入结束节点")
        
        # 检查是否已有交互ID（interrupt返回后的继续执行）
        existing_interaction_id = getattr(self, '_finish_interaction_id', None)
        
        if existing_interaction_id:
            # 继续执行阶段：使用已有交互ID
            logger.info(f"🔄 继续执行finish节点，交互ID：{existing_interaction_id}")
            interaction_id = existing_interaction_id
            user_has_more = interrupt({
                "interaction_type": "human_ask",
                "interaction_id": interaction_id
            })
            # 人机交互后重置迭代计数，给大模型新的处理机会
            state["iteration_count"] = 0
        else:
            # 首次执行阶段：创建交互
            logger.info("🆕 首次执行finish节点，创建人机交互")
            
            # 创建人机交互消息
            interaction_message = self.create_human_interaction_message(
                tool_name="任务完成确认",
                tool_params={},
                interaction_type="human_ask",
                interaction_id=f"finish_confirm_{int(datetime.now().timestamp() * 1000)}"
            )
            
            # 更新交互消息内容
            interaction_message["data"]["query"] = "当前任务已完成，是否还有其他需要补充或处理的问题？"
            interaction_data = interaction_message["data"]
            interaction_id = interaction_data["interaction_id"]
            
            # 保存交互ID到Agent实例属性
            self._finish_interaction_id = interaction_id
            
            # 使用interrupt函数暂停执行，等待用户输入
            user_has_more = interrupt({
                "query": interaction_data["query"],
                "tool_info": interaction_data["tool_info"],
                "interaction_type": "human_ask",
                "interaction_id": interaction_id
            })
        
        # 处理用户响应
        if user_has_more:
            logger.info("🔄 用户有更多问题，继续思考循环")
            
            # 立即清理交互ID，避免循环问题
            if hasattr(self, '_finish_interaction_id'):
                delattr(self, '_finish_interaction_id')
            
            # 获取用户补充消息
            user_message = self._get_user_additional_message(interaction_id)
            
            # 设置新的输入和状态
            if user_message and user_message.strip():
                state["input"] = user_message.strip()
                logger.info(f"🔄 用户补充问题：{user_message}")
                
                # 记录用户补充问题到执行记忆
                if not state.get("this_memory_context"):
                    state["this_memory_context"] = []
                state["this_memory_context"].append({
                    "type": "user_additional_input",
                    "content": user_message,
                    "timestamp": datetime.now().isoformat()
                })
            else:
                state["input"] = "请继续处理我还有其他问题"
                logger.info("🔄 使用默认提示继续处理")
            
            state["should_continue"] = True
            # 用户提供新问题时重置迭代计数，给大模型充分的处理机会
            state["iteration_count"] = 0
            return state
        
        # 用户没有更多问题，正式结束任务
        logger.info("🏁 用户确认没有更多问题，正式结束导航任务")
        
        # 清理交互ID
        if hasattr(self, '_finish_interaction_id'):
            delattr(self, '_finish_interaction_id')
        
        # 设置结束状态
        final_response = "任务已完成。"
        state["final_response"] = final_response
        state["success"] = True
        state["should_continue"] = False

        # 保存执行记忆到数据库
        self._save_execution_memory(state)
        
        logger.info("✅ 导航任务结束")
        return state

    def _get_user_additional_message(self, interaction_id: str) -> str:
        """获取用户补充消息"""
        user_message = None
        
        # 从交互状态中获取用户消息
        if hasattr(self, 'interaction_states') and interaction_id in self.interaction_states:
            interaction_state = self.interaction_states[interaction_id]
            user_message = interaction_state.get("user_message")
            
            if user_message and user_message.strip():
                # 清理交互状态
                del self.interaction_states[interaction_id]
                return user_message.strip()
        
        return user_message or ""

    def _save_execution_memory(self, state: NavigationAgentState) -> None:
        """保存执行记忆到数据库"""
        if not self.memory_manager:
            return
            
        original_input = state.get("original_input", "")
        this_memory_context = state.get("this_memory_context", [])
        final_response = state.get("final_response", "")

        if final_response and original_input:
            try:
                if this_memory_context:
                    memory_summary = self._build_memory_summary(this_memory_context)
                    self.memory_manager.add_message(
                        f" {original_input}",
                        f" {memory_summary}"
                    )
                logger.info("✅ 执行记忆已保存到数据库")
            except Exception as e:
                logger.warning(f"保存记忆到数据库失败: {e}")

    def _decide_next_action(self, state: NavigationAgentState) -> str:
        """决定下一步执行哪个节点"""
        try:
            intelligent_response = state.get("intelligent_response", {})
            action = intelligent_response.get("action", "direct_answer")
            should_continue = intelligent_response.get("should_continue", True)

            # 检查是否应该继续思考
            if not should_continue:
                logger.info("🏁 智能响应指示不继续思考，结束流程")
                state["parsed_action"] = "finish"
                state["should_continue"] = False
                return "finish"

            # 根据智能响应的action类型决定下一步
            if "tool" in action.lower() or "mcp" in action.lower():
                if "tool" in action.lower():
                    state["parsed_action"] = "use_tool"
                else:
                    state["parsed_action"] = "use_mcp"
                next_action = "check_human"
            elif "rag" in action.lower():
                state["parsed_action"] = "use_rag"
                next_action = "rag"
            elif "answer" in action.lower():
                state["parsed_action"] = "direct_answer"
                state["should_continue"] = False
                final_content = intelligent_response.get("content", "任务已完成")
                state["final_response"] = final_content
                next_action = "finish"
            else:
                logger.warning(f"⚠️ 未知的action类型: {action}")
                state["parsed_action"] = "unknown"
                state["should_continue"] = False
                next_action = "finish"

            logger.info(f"智能决策: {action} -> {next_action}")
            return next_action

        except Exception as e:
            logger.error(f"❌ 决策失败: {str(e)}")
            state["parsed_action"] = "direct_answer"
            state["should_continue"] = False
            return "finish"

    def _should_continue_thinking(self, state: NavigationAgentState) -> str:
        """判断是否继续思考循环"""
        return should_continue_execution(state)

    def _should_continue_after_finish(self, state: NavigationAgentState) -> str:
        """判断finish节点后是否继续思考"""
        should_continue = state.get("should_continue", False)
        return "continue" if should_continue else "end"

    def _build_thought_prompt(self, state: NavigationAgentState) -> str:
        """构建思考提示词"""
        input_text = state.get("input", "")
        history_memory_context = state.get("history_memory_context", {})
        this_memory_context = state.get("this_memory_context", [])

        return NavigationPrompts.build_thought_prompt(
            input_text, history_memory_context, this_memory_context
        )

    def _prepare_input_with_memory(self, input_text: str) -> Dict[str, Any]:
        """使用baseAgent的记忆工作流准备输入 - NavigationAgent特有封装"""
        return self._prepare_input_with_memory_base(input_text)

    def _invoke_llm(self, prompt: str):
        """调用LLM 用于生成最终响应"""
        try:
            if not self.llm:
                raise ValueError("LLM未初始化")

            # 创建简单的提示词模板用于最终响应生成
            simple_prompt = ChatPromptTemplate.from_messages([
                SystemMessage(content=self.sys_prompt or "你是一个智能导航助手。"),
                HumanMessage(content=prompt)
            ])

            # 直接调用LLM
            chain = simple_prompt | self.llm
            result = chain.invoke({})

            return result

        except Exception as e:
            logger.error(f"❌ LLM调用失败: {str(e)}")
            # 返回错误响应
            return type('LLMResponse', (), {'content': f"LLM调用失败: {str(e)}"})()

    def _check_human_node(self, state: NavigationAgentState) -> NavigationAgentState:
        """检查是否需要人机交互的节点"""
        logger.info("🔍 检查是否需要人机交互")

        try:
            intelligent_response = state.get("intelligent_response", {})
            parsed_action = intelligent_response.get("action")
            state["parsed_action"] = parsed_action
            state["need_human_interaction"] = False

            if "human" in parsed_action:
                logger.info("🤝 需要人机交互")
                state["need_human_interaction"] = True
            elif "tool" in parsed_action:
                tool_name = intelligent_response.get("tool_name", "")
                if self.check_tool_needs_approval(tool_name):
                    logger.info(f"🤝 工具 {tool_name} 需要人机交互")
                    state["need_human_interaction"] = True
            elif "mcp" in parsed_action:
                server_name = intelligent_response.get("server_name", "")
                tool_name = intelligent_response.get("tool_name", "")
                if self.check_mcp_needs_approval(server_name, tool_name):
                    logger.info(f"🤝 MCP工具 {server_name}:{tool_name} 需要人机交互")
                    state["need_human_interaction"] = True

        except Exception as e:
            logger.error(f"❌ 检查人机交互失败: {str(e)}")
            state["need_human_interaction"] = False

        return state

    def _interrupt_human_node(self, state: NavigationAgentState) -> NavigationAgentState:
        """人机交互节点"""
        logger.info("🤝 进入人机交互节点")
        return self.create_interrupt_and_update_state(state, "nav_interaction")

    def _decide_after_human_check(self, state: NavigationAgentState) -> str:
        """在人机交互检查后决定下一步"""
        try:
            need_interaction = state.get("need_human_interaction", False)
            parsed_action = state.get("parsed_action", "")

            if need_interaction:
                return "human"
            else:
                if "tool" in parsed_action:
                    return "tools"
                elif "mcp" in parsed_action:
                    return "mcp"
                else:
                    return "tools"

        except Exception as e:
            logger.error(f"❌ 人机交互检查后决策失败: {str(e)}")
            return "tools"

    def _decide_after_human_interaction(self, state: NavigationAgentState) -> str:
        """在人机交互后决定下一步"""
        try:
            user_confirmation = state.get("user_confirmation", False)
            parsed_action = state.get("parsed_action", "")

            if not user_confirmation:
                logger.info("❌ 用户取消操作，结束流程")
                state["final_response"] = "用户取消了操作执行"
                state["should_continue"] = False
                return "finish"

            # 用户确认，继续执行
            if "tool" in parsed_action:
                return "tools"
            elif "mcp" in parsed_action:
                return "mcp"
            else:
                return "finish"

        except Exception as e:
            logger.error(f"❌ 人机交互后决策失败: {str(e)}")
            return "finish"

    async def _send_node_update(self, node_name: str, node_result: Dict[str, Any]) -> Dict[str, Any]:
        """发送节点更新消息 - 重写父类方法以提供NavigationAgent特有的高级功能"""
        node_messages = {
            "think": "🤔",
            "execute_tools": "🔧 正在调用工具...",
            "execute_mcp": "🔌 正在调用MCP服务...",
            "execute_rag": "📚 正在调用RAG检索...",
        }

        base_message = node_messages.get(node_name, f"正在执行 {node_name}...")

        # 准备参数
        tool_name = None
        server_name = None
        query = None
        count = 0

        if node_name == "think" and node_result:
            # 对于思考节点，直接使用 current_thought 的内容
            think_content = node_result.get("current_thought", "")
            if think_content:
                base_message = think_content
        elif node_name == "execute_tools" and node_result:
            # 从工具结果中提取工具名称
            tool_results = node_result.get("tool_results", [])
            if tool_results:
                latest_tool = tool_results[-1]
                tool_name = latest_tool.get("tool_name", "未知工具")
        elif node_name == "execute_mcp" and node_result:
            # 从MCP结果中提取服务和工具名称
            mcp_results = node_result.get("mcp_results", [])
            if mcp_results:
                latest_mcp = mcp_results[-1]
                server_name = latest_mcp.get("server_name", "未知服务")
                tool_name = latest_mcp.get("tool_name", "未知工具")
        elif node_name == "execute_rag" and node_result:
            # 从RAG结果中提取查询信息
            rag_results = node_result.get("rag_results", [])
            query = node_result.get("input", "未知查询")
            count = len(rag_results) if rag_results else 0

        # 使用SendResponseDTO创建节点更新消息
        return SendResponseDTO.create_node_update_message(
            node_type=node_name,
            base_content=base_message,
            node_result=node_result,
            stage=f"node_{node_name}",
            tool_name=tool_name,
            server_name=server_name,
            query=query,
            count=count
        )

    def _save_result_to_memory(self, result_type: str, result_data: Dict[str, Any],
                               state: NavigationAgentState) -> None:
        """将执行结果保存到当前执行记忆中"""
        try:
            success = result_data.get("success", False)

            # 构建执行记忆对象
            execute_memory = {
                "type": "execute",
                "execute_type": result_type,  # "tool", "mcp", "rag"
                "success": success,
                "content": result_data.copy(),  # 保存完整的结果数据
                "timestamp": datetime.now().isoformat()
            }

            # 保存到当前执行记忆序列中
            state["this_memory_context"] = state.get("this_memory_context", [])
            state["this_memory_context"].append(execute_memory)

            logger.debug(f"✅ {result_type} 结果已保存到当前执行记忆")

        except Exception as e:
            logger.warning(f"保存{result_type}结果到当前执行记忆失败: {e}")

    def _build_memory_summary(self, memory_context: List[Dict[str, Any]]) -> str:
        """构建执行记忆的详细摘要用于保存到数据库"""
        try:
            summary_parts = []

            for i, memory in enumerate(memory_context, 1):
                memory_type = memory.get("type", "unknown")
                timestamp = memory.get("timestamp", "")

                if memory_type == "think":
                    content = memory.get("content", {})
                    reasoning = content.get("reasoning", "")
                    action = content.get("action", "")
                    tool_name = content.get("tool_name", "")
                    server_name = content.get("server_name", "")

                    think_detail = f"{i}.🤔思考: {action}"
                    if tool_name:
                        think_detail += f" [工具:{tool_name}]"
                    if server_name:
                        think_detail += f" [MCP:{server_name}]"
                    think_detail += f" - {reasoning}"
                    summary_parts.append(think_detail)

                elif memory_type == "execute":
                    execute_type = memory.get("execute_type", "")
                    success = memory.get("success", False)
                    status = "✅成功" if success else "❌失败"
                    content = memory.get("content", {})

                    if execute_type == "tool":
                        tool_name = content.get("tool_name", "")
                        params = content.get("params", {})
                        result = content.get("result", content.get("error", ""))

                        # 构建详细的工具执行信息
                        tool_detail = f"{i}.🔧执行工具: {tool_name} {status}"

                        # 添加参数信息
                        if params:
                            param_strs = []
                            for k, v in params.items():
                                if len(str(v)) > 50:
                                    param_strs.append(f"{k}={str(v)[:50]}...")
                                else:
                                    param_strs.append(f"{k}={v}")
                            tool_detail += f" | 参数:[{', '.join(param_strs)}]"

                        # 添加结果信息
                        if result:
                            result_str = str(result)[:200] + "..." if len(str(result)) > 200 else str(result)
                            tool_detail += f" | 结果:{result_str}"

                        summary_parts.append(tool_detail)

                    elif execute_type == "mcp":
                        server_name = content.get("server_name", "")
                        tool_name = content.get("tool_name", "")
                        params = content.get("params", {})
                        result = content.get("result", content.get("error", ""))

                        # 构建详细的MCP执行信息
                        mcp_detail = f"{i}.🔌执行MCP: {server_name}:{tool_name} {status}"

                        # 添加参数信息
                        if params:
                            param_strs = []
                            for k, v in params.items():
                                if len(str(v)) > 50:
                                    param_strs.append(f"{k}={str(v)[:50]}...")
                                else:
                                    param_strs.append(f"{k}={v}")
                            mcp_detail += f" | 参数:[{', '.join(param_strs)}]"

                        # 添加结果信息
                        if result:
                            result_str = str(result)[:200] + "..." if len(str(result)) > 200 else str(result)
                            mcp_detail += f" | 结果:{result_str}"

                        summary_parts.append(mcp_detail)

                    elif execute_type == "rag":
                        query = content.get("query", "")
                        results_count = content.get("results_count", 0)
                        results = content.get("results", [])

                        # 构建详细的RAG查询信息
                        rag_detail = f"{i}.📚执行查询: '{query}' {status} | 找到{results_count}个结果"

                        # 添加查询结果详情
                        if results and len(results) > 0:
                            rag_detail += " | 结果详情:["
                            result_details = []
                            for j, result in enumerate(results[:3], 1):  # 只显示前3个结果
                                title = result.get("title", "无标题")
                                content_text = result.get("content", "")[:100]
                                url = result.get("url", "")
                                result_details.append(f"{j}.{title}({url}): {content_text}...")
                            rag_detail += " | ".join(result_details)
                            if len(results) > 3:
                                rag_detail += f" | 还有{len(results) - 3}个结果"
                            rag_detail += "]"

                        summary_parts.append(rag_detail)

                elif memory_type == "user_additional_input":
                    # 用户补充输入
                    content_text = memory.get("content", "")
                    user_input_detail = f"{i}.💬用户补充问题: {content_text}"
                    summary_parts.append(user_input_detail)

            return "\n".join(summary_parts)

        except Exception as e:
            logger.warning(f"构建详细记忆摘要失败: {e}")
            return f"执行记忆序列包含{len(memory_context)}个步骤，详情获取失败: {str(e)}"
